import pickle
import numpy as np
from lilab.OpenLabCluster_train.a1_mirror_mutual_filt_clippredpkl import filt_mirror_df
import matplotlib.pyplot as plt
import os.path as osp
import argparse
import pandas as pd


clippredfile_candi = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32_k36/semiseq2seq_iter1/output/semisupervise-decSeq-iter0-epoch5/olc-2024-05-23-semiseq2seq-iter0-epoch5_pca12.clippredpkl'
clippredfile_ref = '/mnt/liying.cibr.ac.cn_Data_Temp/multiview_9/chenxf/00_BehaviorAnalysis-seq2seq/SexAge/Day55_Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/representitive_k36_filt_perc65/Representive_K36.clippredpkl'


def show(output_dir, clippreddata_c):
    embedding_d2 = clippreddata_c['embedding_d2']
    label_pred = clippreddata_c['cluster_labels']
    repr_new_iter_perc = np.mean(clippreddata_c['df_clipNames']['cluster_labels_clean']>0)
    plt.figure(figsize=(12, 10))
    plt.scatter(embedding_d2[::2, 0], embedding_d2[::2, 1], c=label_pred[::2], s=1, cmap='hsv') #
    plt.colorbar()
    plt.title((f'Repr. Proportion {repr_new_iter_perc*100:.2f} % / {len(label_pred)}(all) \n'
                f'Accuracy traintest set to 100 \n'))
    plt.xticks([])
    plt.yticks([])
    plt.savefig(osp.join(output_dir, f'olc-2024-05-23-semiseq2seq.png'))


def main(clippredfile_candi, clippredfile_ref):
    output_dir = osp.dirname(clippredfile_candi)

    clippreddata_c = pickle.load(open(clippredfile_candi, 'rb'))
    clippreddata_r = pickle.load(open(clippredfile_ref, 'rb'))
    df_clipNames_c = clippreddata_c['df_clipNames']
    df_clipNames_c['cluster_labels'] = clippreddata_c['cluster_labels']
    df_clipNames_r = clippreddata_r['df_clipNames']
    nK_mutual = clippreddata_c['nK_mutual']
    nK_mirror_half = clippreddata_c['nK_mirror_half']

    if len(df_clipNames_c) > len(df_clipNames_r):
        df_clipNames_r['cluster_labels'] = clippreddata_r['cluster_labels']
        df_clipNames_r = filt_mirror_df(df_clipNames_r, nK_mutual, nK_mirror_half, start=1)
        df_clipNames_r = df_clipNames_r[df_clipNames_r['cluster_labels_clean']>0]
        df_clipNames_r['cluster_labels_clean2'] = df_clipNames_r['cluster_labels_clean']
        del df_clipNames_r['cluster_labels']
        del df_clipNames_r['cluster_labels_clean']
        df_clipNames_c2 = pd.merge(df_clipNames_c, df_clipNames_r, on=['vnake','isBlack','startFrame'], how='left')
        ind_valid = ~df_clipNames_c2['cluster_labels_clean2'].isna()
        df_clipNames_c2.loc[ind_valid, 'cluster_labels'] = df_clipNames_c2['cluster_labels_clean2'][ind_valid]
        df_clipNames_c['cluster_labels'] = df_clipNames_c2['cluster_labels'].values.astype(int)
        assert (df_clipNames_c['cluster_labels']>0).all()
        df_clipNames_c = filt_mirror_df(df_clipNames_c, nK_mutual, nK_mirror_half, start=1)
    elif len(df_clipNames_c) == len(df_clipNames_r):
        assert (df_clipNames_c[['vnake','isBlack','startFrame']].values ==
                df_clipNames_r[['vnake','isBlack','startFrame']].values).all()
        cluster_labels_r_clean = df_clipNames_r['cluster_labels_clean'].values
        ind_valid = cluster_labels_r_clean > 0
        df_clipNames_c.loc[ind_valid, 'cluster_labels'] = cluster_labels_r_clean[ind_valid]
        df_clipNames_c = filt_mirror_df(df_clipNames_c, nK_mutual, nK_mirror_half, start=1)
    else:
        raise ValueError('len(df_clipNames_c) > len(df_clipNames_r)')

    clippreddata_c['cluster_labels'] = df_clipNames_c['cluster_labels'].values
    clippreddata_c['df_clipNames'] = df_clipNames_c
    pickle.dump(clippreddata_c, open(clippredfile_candi, 'wb'))

    show(output_dir, clippreddata_c)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('clippredfile_candi', type=str)
    parser.add_argument('clippredfile_ref', type=str)
    args = parser.parse_args()
    clippredfile_candi, clippredfile_ref = args.clippredfile_candi, args.clippredfile_ref
    main(clippredfile_candi, clippredfile_ref)
